import pandas as pd
from sklearn.model_selection import train_test_split
import os


def stratified_sample_tsv(input_files, output_files, sample_percentage=0.1, random_seed=42):
    """
    完整版分层抽样TSV文件，保持类别比例

    参数:
        input_files: 输入文件路径列表
        output_files: 输出文件路径列表
        sample_percentage: 抽样比例(0-1之间)
        random_seed: 随机种子，确保可重复性

    返回:
        None (结果会保存到输出文件中)
    """
    # 验证输入
    if len(input_files) != len(output_files):
        raise ValueError("输入文件和输出文件数量不匹配")

    if not 0 < sample_percentage <= 1:
        raise ValueError("抽样比例必须在0到1之间")

    # 创建输出目录（如果不存在）
    for out_file in output_files:
        os.makedirs(os.path.dirname(out_file), exist_ok=True)

    # 处理每个文件
    for in_file, out_file in zip(input_files, output_files):
        print(f"\n正在处理文件: {in_file}")

        try:
            # 读取数据
            df = pd.read_csv(in_file, sep='\t', on_bad_lines='warn')

            # 检查必需的列
            if 'label' not in df.columns:
                raise ValueError(f"文件 {in_file} 中缺少 'label' 列")

            # 计算原始数据分布
            original_dist = df['label'].value_counts(normalize=True).sort_index()
            print("原始数据类别分布:")
            print(original_dist.to_string())

            # 分层抽样
            if len(df) * sample_percentage < 1:
                print("警告：抽样比例太小，至少保留1条数据")
                sample_size = 1
            else:
                sample_size = sample_percentage

            sampled_df, _ = train_test_split(
                df,
                train_size=sample_size,
                stratify=df['label'],
                random_state=random_seed
            )

            # 计算抽样后分布
            sampled_dist = sampled_df['label'].value_counts(normalize=True).sort_index()
            print("\n抽样后类别分布:")
            print(sampled_dist.to_string())

            # 保存结果
            sampled_df.to_csv(out_file, sep='\t', index=False)
            print(f"\n已保存 {len(sampled_df)} 条数据(原始数据的{sample_percentage * 100:.1f}%)到 {out_file}")

        except Exception as e:
            print(f"处理文件 {in_file} 时出错: {str(e)}")
            continue


if __name__ == "__main__":
    # 文件配置
    input_files = [
        'weibo_tsv/train.tsv',
        'weibo_tsv/dev.tsv',
        'weibo_tsv/test.tsv'
    ]

    output_files = [
        'weibo_mini_tsv/train.tsv',
        'weibo_mini_tsv/dev.tsv',
        'weibo_mini_tsv/test.tsv'
    ]

    # 抽样参数
    SAMPLE_PERCENTAGE = 0.1  # 10%
    RANDOM_SEED = 42

    # 执行抽样
    print("=" * 50)
    print(f"开始分层抽样(比例: {SAMPLE_PERCENTAGE * 100}%, 随机种子: {RANDOM_SEED})")
    stratified_sample_tsv(input_files, output_files, SAMPLE_PERCENTAGE, RANDOM_SEED)
    print("\n抽样完成!")